# -*- coding: utf-8 -*-
from universal.common_header import *
from universal.common_header_widget import *
from multiprocessing import Process
from Module.xRayDetection.extra_page import QComboCheckBox
import pyqtgraph as pg
from Module.TrainAndTest.TrainAndTest_Paint import PaintArea
from DeepLearning.project.util_my import main, get_time, get_configs
from mylibs.XmlWriter import SaveProjectToXml
from universal.widget_common import Table_widget,Table_widget_setItem,\
    lineEdit_act_file,createLabel,SetCellWidget,SetReadOnly
from multiprocessing import Process, Queue

FILE_PATH = dir_root()

class trainWidget(QWidget):
    def __init__(self, parent=None):
        super(trainWidget, self).__init__(parent)
        self.leftwidget = QWidget()
        self.trainRightContainer = QWidget()
        # self.leftwidget.setPalette(QPalette(QColor(61,161,91)))
        # self.leftwidget.setAutoFillBackground(True) #要结合调色器生成
        # self.leftwidget.setPalette(Yolo3_trainSettings.createPalette(self,QColor(252,255,255)))
        #self.trainLeftLayout.addSpacing(20)  #在布局末尾加个空白
        self.trainMainLayout = QHBoxLayout(self)
        self.trainLeftLayout = QVBoxLayout(self.leftwidget)
        self.trainRightLayout = QVBoxLayout(self.trainRightContainer)

        self.Tab()  #self.TabelParameter_Input
        self.Show_Result()  #self.TabelParameter_Show
        self.label = PaintArea()
        self.trainPicturesWidget = Demo()
        self.trainPicturesWidget.Message = self.TextBrowser
        #----------------------------
        self.trainLeftLayout.addWidget(self.TabWidget)
        self.trainLeftLayout.addWidget(self.TabelParameter_Show)
        self.trainRightLayout.addWidget(self.label)
        self.trainRightLayout.addWidget(self.trainPicturesWidget)
        self.trainMainLayout.addWidget(self.trainRightContainer)
        # self.trainMainLayout.setStretchFactor(trainDock,0.24)  #设置不可缩放
        # self.trainMainLayout.setStretchFactor(trainRightLayoutContainer,1)
        # self.setLayout(self.trainMainLayout)
        self.init_parameter()

    def init_parameter(self):
        self.labelImgsPath = ''
        self.od_cfg = {}
        self.device_list = []
        self.setPalette(QPalette(QColor(61, 161, 191)))
        self.setAutoFillBackground(True)
        self.TrainModel = ''  # 模型名字
        self.LoadPreTrainingWeight = ''
        self.InputImagesSize = '640*640'
        self.Epochs_input = '2000'
        self.BatchSize = '8'
        self.AlreadyTime_txt = ''
        self.loss = ''
        self.epochEnd = ''
        self.Precision = ''
        self.Recall = ''
        self.mAP50 = ''
        self.Weights_path = ''
        # self.main = queue.Queue()
        self.main = Queue()
        self.projectTask = ''    #目标检测 或者 语义分割 等等...
        self.labelClasses = []
        self.TrainDate = ''
        self.ini_ = 1
        self.refreshTrainInterface()

    def Tab(self):
        self.TabWidget = QTabWidget()
        self.Input_Value()
        self.Input_Value_Senior()
        self.TextBrowser_()
        self.TabWidget.addTab(self.TabelParameter_Input, "基础参数")
        self.TabWidget.addTab(self.TabelParameter_Senior, "高级参数")
        self.TabWidget.addTab(self.TextBrowser, "运行信息")

    def Input_Value(self):
        self.TabelParameter_Input = Table_widget(self, 7, 2)
        name_list = ["模型选择", "预训练模型", "输入图像尺寸", 'Device',"Epochs",
                     "Batch Size", '开始&关闭']
        Table_widget_setItem(self.TabelParameter_Input, name_list)

        self.modelComBox = QComboBox(self) #1
        self.weightpathLineEdit = lineEdit_act_file(self, "Pt files (*.pt)")
        self.imgSizeComBox = QComboBox(self)#3
        self.imgSizeComBox.setFont(QFont(FONTSTR, FONT_SIZE3))
        self.lineEditbox = QLineEdit()
        self.lineEditbox.setFont(QFont(FONTSTR, FONT_SIZE3))
        self.lineEditbox.setAlignment(Qt.AlignCenter)
        self.imgSizeComBox.setLineEdit(self.lineEditbox)
        items = ['320*320', '416*416', '640*640', '800*800', '1280*1280']
        self.imgSizeComBox.clear()
        self.imgSizeComBox.addItems(items)
        self.imgSizeComBox.setCurrentIndex(-1)
        self.selectDevice = QComboCheckBox()

        self.epochsLineEdit = QLineEdit('5000')
        self.batchsizeLineEdit = QLineEdit('4')
        self.startTrainBtn = QPushButton('Start Training')
        self.startTrainBtn.setStyleSheet("background-color:rgb(0,125,0)")
        self.startTrainBtn.setAutoFillBackground(True)
        self.startTrainBtn.setFont(QFont(FONTSTR, FONT_SIZE3))
        self.startTrainBtn.setEnabled(True)
        self.startTrainBtn.clicked.connect(self.startTrain)

        TabelParameter_Input_list = [self.modelComBox, self.weightpathLineEdit, self.imgSizeComBox,self.selectDevice,
              self.epochsLineEdit, self.batchsizeLineEdit , self.startTrainBtn]
        SetCellWidget(self.TabelParameter_Input,TabelParameter_Input_list)
        self.TabelParameter_Input.setContentsMargins(0, 0, 0, 0)

    def Input_Value_Senior(self):
        self.TabelParameter_Senior = Table_widget(self, 2, 2)
        name_list = ["Validate Split", 'Learning Rate']
        Table_widget_setItem(self.TabelParameter_Senior, name_list)
        self.TabelParameter_Senior.setContentsMargins(0, 0, 0, 0)

    def TextBrowser_(self):
        self.TextBrowser = QTextBrowser()
        self.TextBrowser.document().setMaximumBlockCount(100)
        self.TextBrowser.append('------信息框------')

    def Show_Result(self):
        self.TabelParameter_Show = Table_widget(self, 7, 2)
        name_list = ['已运行时间',"剩余时间", "当前损失","当前轮数", "precision", 'recall', 'mAP' ]
        Table_widget_setItem(self.TabelParameter_Show, name_list)#
        self.Already_Time = QLineEdit()
        self.Remain_Time = QLineEdit()
        self.CurrentLoss_Label = QLineEdit()
        self.CurrentEpoch_Label = QLineEdit()
        self.Precision_Label = QLineEdit()
        self.Recall_Label = QLineEdit()
        self.Map_Label = QLineEdit()

        TabelParameter_Show_list = [self.Already_Time,self.Remain_Time,self.CurrentLoss_Label,
              self.CurrentEpoch_Label,self.Precision_Label,self.Recall_Label,self.Map_Label]
        SetReadOnly(TabelParameter_Show_list)
        SetCellWidget(self.TabelParameter_Show, TabelParameter_Show_list)
        self.TabelParameter_Show.setContentsMargins(0, 0, 0, 0)

    def AddModelSelect(self):
        try:
            self.modelComBox.clear()
            if self.projectTask == '目标检测':
                self.englishTask = 'detection'
                self.model_False = ["OD_n","OD_n2","OD_s","OD_s2","OD_m","OD_m2","OD_l","OD_l2","OD_x","OD_x2"]
            if self.projectTask == '图像分类':
                self.englishTask = 'classification'
                self.model_False = ["OD_n_cls", "OD_s_cls", "OD_m_cls", "OD_l_cls", "OD_x_cls", "OD_e_cls_0", "OD_e_cls_1",
                                    "OD_e_cls_2", "OD_e_cls_3", "OD_r_cls_0", "OD_r_cls_1", "OD_r_cls_2", "OD_r_cls_3"]
            if self.projectTask == '语义分割':
                self.englishTask = 'detection'
            if self.projectTask == '实例分割':
                self.englishTask = 'segmentation'
                self.model_False = ["OD_n_seg","OD_s_seg","OD_m_seg","OD_l_seg","OD_x_seg"]

            self.model_true = get_configs(self.englishTask)
            for i in self.model_False:
                self.modelComBox.addItem(str(i))
        except Exception as e:
            self.TextBrowser.append('Function:AddModelSelect')
            self.TextBrowser.append(str(e))

    def TrainsetText(self):
        self.startTrainBtn.setText('开始')
        self.startTrainBtn.setStyleSheet('''QPushButton{background:rgb(0,125,0);
            border-radius:5px;}QPushButton:hover{background:blue;}''')

    def startTrain(self):
        try:
            if self.ini_ == 1:
                self.ini_ = 0
            else:
                self.ini_ = 1
            if self.ini_ == 1:
                self.main.put({'stop': True})
                self.TextBrowser.append('训练结束！')
                self.TrainsetText()
                self.SaveSetting()
                return
            else:
                # self.main.queue.clear()
                self.trainPicturesWidget.PlotPoint()
                self.TextBrowser.append('训练开始！')
                self.startTrainBtn.setText('停止')
                self.startTrainBtn.setStyleSheet('''QPushButton{background:rgb(255,0,0);
                    border-radius:5px;}QPushButton:hover{background:blue;}''')
                self.AlreadyTime = 0
                kwargs,sub = self.Load_Parameter()

                self.UPloadThread = Process(target=main, kwargs=kwargs)
                # self.UPloadThread.daemon = True
                self.UPloadThread.start()

                self.GetValThread = threading.Thread(target=self.GetVal, args=(self.main,sub))
                self.GetValThread.daemon = True
                self.GetValThread.start()
        except Exception as e:
            self.TextBrowser.append('Function:startTrain')
            self.TextBrowser.append(str(e))

    def SaveSetting(self):
        try:
            self.TrainDate = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))
            self.od_cfg["trainModel"] = self.modelComBox.currentText().strip()
            self.od_cfg["trainPreWeightpath"] = self.weightpathLineEdit.text().strip()
            self.od_cfg["trainInputImageSize"] = self.imgSizeComBox.currentText().strip()
            self.od_cfg["epochs"] = self.epochsLineEdit.text()
            self.od_cfg["batchSize"] = self.batchsizeLineEdit.text()
            self.od_cfg["traintime"] = self.AlreadyTime_txt
            self.od_cfg["trainValloss"] = self.loss
            self.od_cfg["trainendEpochs"] = self.epochEnd
            self.od_cfg["trainDate"] = self.TrainDate
            self.od_cfg["precision"] = self.Precision
            self.od_cfg["recall"] = self.Recall
            self.od_cfg["mAP"] = self.mAP50
            SaveProjectToXml(self.od_cfg,None)

        except Exception as e:
            self.TextBrowser.append('Function:SaveSetting')
            self.TextBrowser.append(str(e))

    def Load_Parameter(self):
        try:
            DeviceAll = self.selectDevice.currentText()  # 选择设备
            if "c" in DeviceAll or 'C' in DeviceAll:
                device = 'cpu'
            else:
                device_list = re.findall(r"\d+\.?\d*", DeviceAll)
                device = str(",".join(device_list))
            # sub = queue.Queue()
            sub = Queue()
            self.modelComBox_text = self.modelComBox.currentText()
            if self.projectTask == '图像分类':
                self.modelComBox_text = self.input_cls(self.modelComBox_text)
            else:
                self.modelComBox_text = self.model_true[self.model_False.index(self.modelComBox_text)]
            self.LoadPreTrainingWeight = self.weightpathLineEdit.text()
            self.InputImagesSize = int(self.imgSizeComBox.currentText().split('*')[0])
            # self.ValidateSplit =self.valSplitLineEdit.text()
            self.BatchSize = self.batchsizeLineEdit.text()
            self.Epochs_input = self.epochsLineEdit.text()
            # index_class = self.labelClasses.index("All")
            # if len(self.labelClasses) > 1 and 0 == index_class:
            #     del self.labelClasses[0]
            cfg = {
                'action': 'train',  # train infer export results configs refactor
                'task': self.englishTask,  # classification detection segmentation
                'data_dir': self.labelImgsPath,  # dataset directory  不能有images
                'cfg_id': self.modelComBox_text,  # config id
                'imsz': self.InputImagesSize,
                'names': self.labelClasses,

                'device': device,  # cpu or 0 or 0,1...
                'out_dir': self.Weights_path,  # save action results
                'weights': self.LoadPreTrainingWeight,  # weights path
                'batch_size': self.BatchSize,
                'epochs': self.Epochs_input,
                'weights_dir': dir_root()+os.sep+'Pre_model',
                'github_dir': os.path.join(dir_root(), 'DeepLearning', '.github'),
                'exist': False,

                'sub_queue': sub,  # queue for sending data from sub thread to main thread, like "progress"
                'main_queue': self.main,  # queue for sending data from main thread to sub thread, like "stop"
                'logger': init_log(),  # custom logger object for logging
            }
            return cfg,sub

        except Exception as e:
            self.TextBrowser.append('Function:Load_Parameter')
            self.TextBrowser.append(str(e))

    def GetVal(self,main,sub):
        while True:
            time.sleep(0.1)
            try:
                flags = main.get(block=False)
            except:
                flags = {}
            if flags.get('stop', False):
                main.put({'stop': True})
                break
            try:
                result = sub.get(block=False)
                if 'err' in result:
                    self.TextBrowser.append(result['err'])
                    self.TrainsetText()
                    self.TextBrowser.append('训练结束！')
                    self.ini_ = 1
                    break
                if 'epoch' in result:
                    train_result = result
                else:
                    val_result = result
                    self.epochEnd = train_result['epoch']
                    OnceTime = train_result['train_time'] + val_result['val_time']
                    Alltime = OnceTime * int(self.Epochs_input)
                    self.AlreadyTime += OnceTime
                    self.AlreadyTime_txt = get_time(self.AlreadyTime)
                    Remaintime = Alltime - self.AlreadyTime
                    self.Remaintime_txt = get_time(Remaintime)

                    if self.englishTask == 'classification':
                        self.loss = round(train_result['vloss'] + train_result['tloss'], 4)
                        self.acc = round(val_result['top1'], 4)
                    elif self.englishTask == 'detection':
                        self.loss = round(train_result['train/box_loss'] + train_result['train/obj_loss']
                                    + train_result['train/cls_loss'], 4)
                        self.acc = round(val_result['map50'], 4)
                        self.Precision = round(val_result['precision'], 4)
                        self.Recall = round(val_result['recall'], 4)
                    elif self.englishTask == 'segmentation':
                        self.loss = round(train_result['train/box_loss'] + train_result['train/obj_loss']
                                    + train_result['train/cls_loss'] + train_result['train/seg_loss'], 4)
                        self.acc = round((val_result['map50_bbox'] + val_result['map50_mask'])/2, 4)
                        self.Precision = round((val_result['mp_bbox'] + val_result['mp_mask'])/2, 4)
                        self.Recall = round((val_result['mr_bbox'] + val_result['mr_mask'])/2, 4)

                    self.trainPicturesWidget.x_int = float(self.epochEnd) #传值
                    self.trainPicturesWidget.y_int_loss = float(self.loss)
                    self.trainPicturesWidget.y_int_map = float(self.acc)
                    self.trainPicturesWidget.plot_slot() #画图

                    self.ShowTrainResult()
                    self.Remain_Time.setText(str(self.Remaintime_txt))
                    self.Map_Label.setText(str(self.acc))
                    self.TextBrowser.append('正常训练完成%d个Epoch'%(train_result['epoch']+1))
                    if train_result['epoch']+1 == int(self.Epochs_input):
                        self.SaveSetting()
                        self.Remain_Time.setText('0')
                        self.TrainsetText()
                        self.TextBrowser.append('训练结束！')
                        self.ini_ = 1
                        break
            except Exception as e:
                if e.args != ():
                    self.TextBrowser.append('Function:GetVal')
                    self.TextBrowser.append(str(e))

    def clearShow(self):
        self.modelComBox.clear()
        self.trainPicturesWidget.PlotPoint()
        self.trainPicturesWidget.set_plot_slot0()
        self.Already_Time.clear()
        self.CurrentLoss_Label.clear()
        self.CurrentEpoch_Label.clear()
        self.Precision_Label.clear()
        self.Recall_Label.clear()
        self.label.label2.clear()

    def ShowTrainResult(self):
        self.Already_Time.setText(str(self.AlreadyTime_txt))
        self.CurrentLoss_Label.setText(str(self.loss))
        self.CurrentEpoch_Label.setText(str(self.epochEnd))
        self.Precision_Label.setText(str(self.Precision))
        self.Recall_Label.setText(str(self.Recall))

    def refreshTrainInterface(self,count=0):#
        try:
            self.ShowTrainResult()
            self.modelComBox.setCurrentText(self.TrainModel)
            self.weightpathLineEdit.setText(self.LoadPreTrainingWeight)
            self.imgSizeComBox.setCurrentText(self.InputImagesSize)
            self.epochsLineEdit.setText(self.Epochs_input)
            self.batchsizeLineEdit.setText(self.BatchSize)
            self.Map_Label.setText(str(self.mAP50))
            if self.TrainDate != "":
                self.label.label2.setText('Date: '+ self.TrainDate)
            if count > 0:
                self.label.label1.setText(str(count) + ' Images')
        except Exception as e:
                self.TextBrowser.append('Function:refreshTrainInterface')
                self.TextBrowser.append(str(e))

    def getTrainValue(self, para_Cfg):
        try:
            self.TrainModel = para_Cfg['trainModel']
            self.LoadPreTrainingWeight = para_Cfg['trainPreWeightpath']
            self.InputImagesSize = para_Cfg['trainInputImageSize']
            self.Epochs_input = para_Cfg['epochs']
            self.BatchSize = para_Cfg['batchSize']
            self.Weights_path = para_Cfg['trainSaveWeigth']
            self.AlreadyTime_txt = para_Cfg['traintime']
            self.loss = para_Cfg['trainValloss']
            self.epochEnd = para_Cfg['trainendEpochs']
            self.TrainDate = para_Cfg['trainDate']
            self.Precision = para_Cfg['precision']
            self.Recall = para_Cfg['recall']
            self.mAP50 = para_Cfg['mAP']
            self.labelClasses = para_Cfg['labelClasses']
            self.refreshTrainInterface()
        except Exception as e:
                self.TextBrowser.append('Function:refreshTrainInterface')
                self.TextBrowser.append(str(e))

    # def imgSizeCobBoxTextChanged(self):
    #     try:
    #         text = self.imgSizeComBox.lineEdit().text()
    #         if '*' in text:
    #             self.imageSize = (int(text.split('*')[0]), int(text.split('*')[0]))
    #     except Exception as e:
    #             self.TextBrowser.append('Function:imgSizeCobBoxTextChanged')
    #             self.TextBrowser.append(str(e))

    def createPalette(self, qcolor):
        try:
            qpalette = QPalette()
            qpalette.setColor(QPalette.Background,qcolor)
            return qpalette
        except Exception as e:
            self.TextBrowser.append('Function:createPalette')
            self.TextBrowser.append(str(e))

    def input_cls(self,model):
        if model == "OD_n_cls":
            return "yolov5n-cls"

        elif model == "OD_s_cls":
            return "yolov5s-cls"

        elif model == "OD_m_cls":
            return "yolov5m-cls"

        elif model == "OD_l_cls":
            return "yolov5l-cls"

        elif model == "OD_x_cls":
            return "yolov5x-cls"

        elif model == "OD_e_cls_0":
            return "yolov5-efficientnet_b0"

        elif model == "OD_e_cls_1":
            return "yolov5-efficientnet_b1"

        elif model == "OD_e_cls_2":
            return "yolov5-efficientnet_b2"

        elif model == "OD_e_cls_3":
            return "yolov5-efficientnet_b3"

        elif model == "OD_r_cls_0":
            return "yolov5-resnet18"

        elif model == "OD_r_cls_1":
            return "yolov5-resnet34"

        elif model == "OD_r_cls_2":
            return "yolov5-resnet50"

        elif model == "OD_r_cls_3":
            return "yolov5-resnet101"
        else:
            return "yolov5n-cls"

class Demo(QWidget):
    def __init__(self):
        super(Demo, self).__init__()
        self.resize(600, 600)
        pg.setConfigOptions(leftButtonPan=True)
        pg.setConfigOptions(antialias=True)     #设置平滑绘制QPainter
        pg.setConfigOption('background', 'w')   #设置背景颜色（包括文本，线条，坐标轴等）  w为白色
        pg.setConfigOption('foreground', 'k')   #设置前景颜色（包括文本，线条，坐标轴等）  k为黑色
        r_symbol = random.choice(['o'])  #选择坐标点形状
        r_symbol2 = random.choice(['h'])  #选择坐标点形状
        # r_color = random.choice(['b', 'g', 'r', 'c', 'm', 'y', 'k', 'd', 'l', 's']) #颜色
        self.PlotPoint()
        self.Message = None
        self.pw = pg.PlotWidget(self)
        self.pw.addLegend()
        chartPen1 = pg.mkPen(color=(192, 80, 77), width=3)
        chartPen2 = pg.mkPen(color=(142, 130, 30), width=3)
        self.plot_data1 = self.pw.plot(self.x_Int,  pen=chartPen1, symbol=r_symbol, symbolBrush='g', name='LOSS Line')
        self.plot_data2 = self.pw.plot(self.x_Int,  pen=chartPen2, symbol=r_symbol2, symbolBrush='k', name='mAP Line')

        self.v_layout = QVBoxLayout()
        self.v_layout.addWidget(self.pw)
        self.setLayout(self.v_layout)

        # 设置坐标范围
        self.pw.setYRange(min=-0.2 , # 最小值
                          max=1)  # 最大值
        # 显示表格线
        self.pw.showGrid(x=True, y=True)

        # 设置上下左右的label
        self.pw.setLabel("left", "value")
        self.pw.setLabel("bottom", "epochs")

        # 设置图表标题
        self.pw.setTitle("训练折线图",
                         size='12pt')

    def PlotPoint(self):
        self.xint = ''
        self.x_int = ''
        self.y_int_loss = ''
        self.y_int_map = ''
        self.x_Int = []  # 存放坐标点的数组
        self.y_Int = []
        self.y_Int_map = []

    # def timer_init(self):
    #     self.timer = QTimer(self)
    #     self.timer.timeout.connect(self.plot_slot)

    def set_plot_slot0(self):
        self.plot_data1.setData()
        self.plot_data2.setData()

    def plot_slot(self):
        try:
            if self.xint != self.x_int:
                self.xint = self.x_int
                if self.x_int != '' and self.y_int_loss != '':
                    print('画图1')
                    self.x_Int.append(float(self.x_int))
                    self.y_Int.append(float(self.y_int_loss))
                    self.plot_data1.setData(self.x_Int,self.y_Int)

                if self.x_int != '' and self.y_int_map != '':
                    print('画图2')
                    self.y_Int_map.append(float(self.y_int_map))
                    r_symbol = random.choice(['o', 's', 't', 't1', 't2', 't3', 'd', '+', 'x', 'p', 'h', 'star'])
                    self.plot_data2.setData(self.x_Int,self.y_Int_map)

        except Exception as e:
            self.Message.append('Function:plot_slot')
            self.Message.append(str(e))

# if __name__ == '__main__':
#     app = QApplication(sys.argv)
#     w = trainWidget()
#     # w.showMaximized()
#     w.show()
#     sys.exit(app.exec())
